#include "mpi/partition_3d.hpp"

void argsParser(int argc, char *argv[], int &p, int &M, int &N, int &K, int &mp,
                int &np, int &kp) {
  if (argc < 5) {
    INFO_PRINTF("argc = %d\n", argc);
    ERROR_PRINTF("Usage: %s <p ><M> <N> <K>\n", argv[0]);
    exit(1);
  }

  p = std::atoi(argv[1]);
  M = std::atoi(argv[2]);
  N = std::atoi(argv[3]);
  K = std::atoi(argv[4]);

  int rp; // Residual processes
  // TODO: remove Cannon's requirement: 'mod(max(mp,np),min(mp,np))=0'
  calc_3d_decomposition(p, M, N, K, &mp, &np, &kp, &rp);
  if ((mp < 1) || (np < 1) || (kp < 1) || (mp * np * kp > p)) {
    ERROR_PRINTF(
        "Invalid process grid generated: p = %d, mp, np, kp = %d, %d, %d\n", p,
        mp, np, kp);
    exit(1);
  }
  // TODO: support rp != 0
  if (rp) {
    WARNING_PRINTF("Unimplemented case: rp = %d\n", rp);
    exit(1);
  }

  printf("Problem Size m  * n  * k  :\t%-5d * %-5d * %-5d\n", M, N, K);
  printf("Process Grid mp * np * kp :\t%-5d * %-5d * %-5d\n", mp, np, kp);
  if (M % mp != 0 || K % kp != 0 || N % np != 0) {
    ERROR_PRINTF("Invalid problem size: mp, np, kp = %d, %d, %d must excatly "
                 "divide M, K and N\n",
                 mp, np, kp);
    exit(1);
  }
  int mb = M / mp;
  int nb = N / np;
  int kb = K / kp;
  printf("Work cuboid  mb * nb * kb :\t%-5d * %-5d * %-5d\n", mb, nb, kb);
  printf("Process utilization       :\t%.2f %% \n",
         100.0 * (1.0 - (double)rp / (double)p));
  fflush(stdout);
}

int main(int argc, char *argv[]) {
  // Get MPI metadata
  int p;

  int M, K, N;
  int mp, kp, np;
  argsParser(argc, argv, p, M, N, K, mp, np, kp);
  return 0;
}